package com.ai.module.websocket.service.impl;

import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.common.constant.StringConstant;
import com.ai.module.user.domain.req.AiProductReq;
import com.ai.module.user.domain.vo.UserLoginInfo;
import com.ai.module.user.service.PublicTitleService;
import com.ai.module.websocket.service.AIProductWebSocketService;
import com.ai.module.websocket.vo.WebSocketProduct;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.ai.module.chat.common.constant.StringConstant.userROBOT;

/**
 * @author: heJinQin
 * @create: 2024-04-13 12:47
 * @description:
 **/
@Slf4j
@Service
public class AIProductWebSocketServiceImpl implements AIProductWebSocketService {

    private final Map<Integer, WebSocketSession> clients = new ConcurrentHashMap<>();


    private final Map<WebSocketSession,String> sessionUserMap = new ConcurrentHashMap<>();


    //    根据用户 userId @ robotId 为 key
    @Autowired
    private PublicTitleService titleService;

    @Override
    public void handleOpen(WebSocketSession session) {
        // 这个时候就需要在建立 webSocket 时存储的 用户信息了
        Map<String, Object> attributes = session.getAttributes();
        UserLoginInfo clientUserInfo = (UserLoginInfo) attributes.get(StringConstant.SOCKET_USER);
        clients.put(clientUserInfo.getId(), session);
        //
        log.info("新的 websocket 有新连接加入！当前在线人数为 [{}] ",clients.size());

    }

    @Override
    public void handleClose(WebSocketSession session) {
        // 这个时候就需要在建立 webSocket 时存储的 用户信息了
        Map<String, Object> attributes = session.getAttributes();
        UserLoginInfo clientUserInfo = (UserLoginInfo) attributes.get(StringConstant.SOCKET_USER);
        clients.remove(clientUserInfo.getId());

        String userKey = sessionUserMap.remove(session);
        if (userKey !=null ){
            sessionUserMap.remove(userKey);
            userProductSessionMap.remove(userKey);
        }
        ThreadLocalContext.clear();
        log.info("新的 websocket 链接关闭！当前在线人数为 ：{}", clients.size());
    }

    @Override
    public void handleMessage(WebSocketSession session, String message) {
        // 只处理前端传来的文本消息，并且直接丢弃了客户端传来的消息
        log.info("新的 websocket 接收到前端的数据为：{}", message);
        try {
            AiProductReq aiProductReq = JSONObject.parseObject(message, AiProductReq.class);
            aiProductReq.setStream(true);
            Map<String, Object> attributes = session.getAttributes();
            UserLoginInfo clientUserInfo = (UserLoginInfo) attributes.get(StringConstant.SOCKET_USER);
            ThreadLocalContext.setUserLoginInfo(clientUserInfo);
            String robotId = aiProductReq.getRobotId();
            String uuid = clientUserInfo.getUuid();
            String userRobotKey = String.format(userROBOT, uuid, robotId);
            if (! userProductSessionMap.containsKey(userRobotKey)){
                userProductSessionMap.put(userRobotKey,session);
                sessionUserMap.put(session,userRobotKey);
            }

            titleService.productChat(aiProductReq);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void sendMessage(WebSocketSession session, String message) throws IOException {
        this.sendMessage(session, new TextMessage(message));
    }

    @Override
    public void sendMessage(Integer userId, TextMessage message) throws IOException {
        WebSocketSession webSocketSession = clients.get(userId);

        if (webSocketSession.isOpen()) {
            webSocketSession.sendMessage(message);
        }
    }

    @Override
    public void sendMessage(Integer userId, String message) throws IOException {
        this.sendMessage(userId, new TextMessage(message));
    }

    @Override
    public void sendMessage(WebSocketSession session, TextMessage message) throws IOException {
        session.sendMessage(message);
    }

    @Override
    public void broadCast(String message) throws IOException {
        // 广播模式 高并发问题
        clients.values().forEach(session -> {
            if (session.isOpen()) {
                try {
                    session.sendMessage(new TextMessage(message));
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        });
    }

    @Override
    public void broadCast(TextMessage message) throws IOException {
        clients.values().forEach(session -> {
            if (session.isOpen()) {
                try {
                    session.sendMessage(message);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        });
    }

    @Override
    public void handleError(WebSocketSession session, Throwable error) {
        log.error("新的 websocket 接收到的错误信息 error：{}，error：{}, session id  [{}]"
                , error.getMessage()
                , error
                , session.getId()

        );
    }


    @Override
    public void sendMessage(WebSocketSession session, WebSocketProduct productMessage)  {
        String jsonString = JSONObject.toJSONString(productMessage);
       try {
           sendMessage(session,jsonString);
       }catch (IOException exception){
           exception.printStackTrace();
           log.error("新的 websocket 发送信息失败, ex[{}],message[{}]",exception,exception.getMessage());
       }
    }

}

